PyTorch LightningのLightningDataModule

主なメソッド

主なメソッドの呼び出され方

prepare_data

downloadやtokenizeをするためのフック

Lightning ensures the prepare_data() is called only within a single process

It is not recommended to assign state here (e.g. self.x = y)

setup

train, val, test, predictのデータを属性にセットすることが多い

count number of classes

build vocabulary

perform train/val/test splits

create datasets

apply transforms (defined explicitly in your datamodule)

stageが指定されなければNone

すべてのstageを実行する実装が多い

if stage in (None, "fit")

This method expects a stage argument. It is used to separate setup logic for trainer.{fit,validate,test,predict}. If setup is called with stage=None, we assume all stages have been set-up.

setup is called from every process across all the nodes. Setting state here is recommended.

{train,val,test,predict}_dataloader